package knowlwdge_graph.knowledge_retrival.yago_retrival;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.api.client.json.Json;
import common.dao.orient.OrientJDBCDao;
import common.helper.json.JavaObjectToJsonNode;
import common.helper.json.JavaToJsonString;
import knowlwdge_graph.knowledge_etl.yago_etl.SchemaDefinition;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;

/**
 * Created by julianzliu on 4/21/2017.
 */
public class YagoEntityRereival {

    static final org.slf4j.Logger logger = LoggerFactory.getLogger(YagoEntityRereival.class);

    private OrientJDBCDao jdbcDao;
    private Connection connection;


    public YagoEntityRereival(){
        this.jdbcDao = new OrientJDBCDao();
        this.connection = this.jdbcDao.getConn();
    }


    private JsonNode retrievalEntity(String sql){
        //String sql = String.format("select *,out.yago_id, in.yago_id from (select expand(outE()) from V where yago_id='%s')", yago_id);
        logger.info("[retrievalEntity] sql = {}",sql);
        Map<String, Map<String,Set<String>> > triples = new HashMap<>();
        Statement stmt = null;
        ResultSet rs = null;
        JsonNode node = null;
        try {
            stmt = this.connection.createStatement();
            rs = stmt.executeQuery( sql );

            while (rs.next()){
                //System.out.println( rs.toString() );
                String out_id = rs.getString("out");
                String in_id = rs.getString("in");
                String edge = rs.getString("yago_label");
                System.out.println("(out=" + out_id + ",edge=" + edge + " ,in=" + in_id + ")");
                Map<String, Set<String>> edges;
                if(!triples.containsKey(out_id)){
                    edges = new HashMap<>();
                    triples.put(out_id, edges );
                }
                else {
                    edges = triples.get(out_id);
                }
                Set<String> values;
                if(edges.containsKey(edge)){
                    values = edges.get(edge);
                }else{
                    values = new HashSet<>();
                    edges.put(edge, values);
                }
                values.add( in_id );
            }
            node = JavaObjectToJsonNode.ObjectToNode( triples );
            stmt.close();
        }
        catch (Exception e){
            e.printStackTrace();
        }

        return node;
    }

    public JsonNode retrievalEntityById(String yago_id){
        String sql = String.format("select *,out.yago_id, in.yago_id from (select expand(outE()) from V where yago_id='%s')", yago_id);
        return retrievalEntity(sql);
    }

    public JsonNode retrievalEntityByLabel(String yago_label){
        yago_label += " +eng";  //只检索英语文本
        String sql = String.format("select *,out.yago_id, in.yago_id from (select expand(inE()) from YAGO_LABEL_CLASS where yago_label LUCENE \"%s\")", yago_label);
        return retrievalEntity(sql);
    }


    /************************************************************
     *
     */
    public static void main(String[] args){
        String yago_id = "<Natasha>";
        if(args.length > 0);
            yago_id = args[0];
        YagoEntityRereival entityRereival = new YagoEntityRereival();
        JsonNode entity = entityRereival.retrievalEntityById( yago_id );

        System.out.println("\n\n[Entity Result = ]\n" + JavaToJsonString.mapJavaToJson(entity) );
        System.out.println("\n\n[Finished]\n\n");

        String label = "+Obama";
        entity = entityRereival.retrievalEntityByLabel(yago_id);

        System.out.println("\n\n[Entity Result = ]\n" + JavaToJsonString.mapJavaToJson(entity) );
        System.out.println("\n\n[Finished]\n\n");
    }
}
